#include <cstdio>
#include <cstdint>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/mman.h>
#include <turbojpeg.h>
#include "zq_loader.h"

struct pixel
{
    uint8_t r;
    uint8_t g;
    uint8_t b;
};

size_t zq_loader(const char* fname, uint8_t* dst);
size_t zq_loader_mmap(const char* fname, uint8_t** dst);
int zq_decoder(tjhandle handle, const uint8_t* src, size_t jpegSize, uint8_t* dst);
int zq_decoder2(tjhandle handle, const uint8_t* src, size_t jpegSize, uint8_t* dst);
int zq_decoder2_resz(tjhandle handle, const uint8_t* src, size_t jpegSize, uint8_t* dst);
void zq_resize(const uint8_t* src, uint8_t* dst);
void zq_resize2(const uint8_t* src, uint8_t* dst);

int loadJpg_raw(const unsigned tid, zq_queue& queue, counter_mt& counter)
{
    tjhandle handle = tjInitDecompress();
    uint8_t* jpegBuf = static_cast<uint8_t*>(malloc(sizeof(uint8_t) * IMG_ROW * IMG_COL * IMG_CHANEL));
    if (jpegBuf == nullptr)
    {
        std::printf("[loadJpg_raw #%u] Allocating JPEG buffer FAILED\n", tid);
        return -1;
    }
    int nitem = 0;
    while (true)
    {
        message info = queue.get();
        if (info.ptr == nullptr)
        {
            break;
        }

        size_t jpegSize = zq_loader(info.path.c_str(), jpegBuf);
        zq_decoder(handle, jpegBuf, jpegSize, info.ptr);

        counter.count();
        ++nitem;
    }
    tjDestroy(handle);
    free(jpegBuf);
    std::printf("[loadJpg_raw #%u] loaded %d jpgs, now exiting...\n", tid, nitem);
    return nitem;
}

int loadJpg_resize(const unsigned tid, zq_queue& queue, counter_mt& counter)
{
    tjhandle handle = tjInitDecompress();
    uint8_t* jpegBuf = static_cast<uint8_t*>(malloc(sizeof(uint8_t) * IMG_ROW * IMG_COL * IMG_CHANEL));
    uint8_t*  rawBuf = static_cast<uint8_t*>(malloc(sizeof(uint8_t) * IMG_ROW * IMG_COL * IMG_CHANEL));
    if (jpegBuf == nullptr || rawBuf == nullptr)
    {
        std::printf("[loadJpg_resize #%u] Allocating JPEG buffer FAILED\n", tid);
        return -1;
    }
    int nitem = 0;
    while (true)
    {
        message info = queue.get();
        if (info.ptr == nullptr)
        {
            break;
        }

        size_t jpegSize = zq_loader(info.path.c_str(), jpegBuf);
        zq_decoder(handle, jpegBuf, jpegSize, rawBuf);
        zq_resize(rawBuf, info.ptr);

        counter.count();
        ++nitem;
    }
    tjDestroy(handle);
    free(jpegBuf);
    free(rawBuf);
    std::printf("[loadJpg_resize #%u] loaded %d jpgs, now exiting...\n", tid, nitem);
    return nitem;
};


int loadJpg_raw_mmap(const unsigned tid, zq_queue& queue, counter_mt& counter)
{
    tjhandle handle = tjInitDecompress();
    uint8_t* jpegMap = nullptr;
    int nitem = 0;
    while (true)
    {
        message info = queue.get();
        if (info.ptr == nullptr)
        {
            break;
        }

        size_t jpegSize = zq_loader_mmap(info.path.c_str(), &jpegMap);
        zq_decoder(handle, jpegMap, jpegSize, info.ptr);
        munmap(static_cast<void*>(jpegMap), jpegSize);
        jpegMap = nullptr;

        counter.count();
        ++nitem;
    }
    tjDestroy(handle);
    std::printf("[loadJpg_raw_mmap #%u] loaded %d jpgs, now exiting...\n", tid, nitem);
    return nitem;
}


int loadJpg_resize_mmap(const unsigned tid, zq_queue& queue, counter_mt& counter)
{
    tjhandle handle = tjInitDecompress();
    uint8_t* jpegMap = nullptr;
    uint8_t*  rawBuf = static_cast<uint8_t*>(malloc(sizeof(uint8_t) * IMG_ROW * IMG_COL * IMG_CHANEL));
    if (rawBuf == nullptr)
    {
        std::printf("[loadJpg_resize_mmap #%u] Allocating JPEG buffer FAILED\n", tid);
        return -1;
    }
    int nitem = 0;
    while (true)
    {
        message info = queue.get();
        if (info.ptr == nullptr)
        {
            break;
        }

        size_t jpegSize = zq_loader_mmap(info.path.c_str(), &jpegMap);
        zq_decoder(handle, jpegMap, jpegSize, rawBuf);
        munmap(static_cast<void*>(jpegMap), jpegSize);
        jpegMap = nullptr;

        zq_resize(rawBuf, info.ptr);

        counter.count();
        ++nitem;
    }
    tjDestroy(handle);
    free(rawBuf);
    std::printf("[loadJpg_resize_mmap #%u] loaded %d jpgs, now exiting...\n", tid, nitem);
    return nitem;
};


int loadJpg_mmap_dec2_resz(const unsigned tid, zq_queue& queue, counter_mt& counter)
{
    tjhandle handle = tjInitDecompress();
    uint8_t* jpegMap = nullptr;
    uint8_t*  rawBuf = static_cast<uint8_t*>(malloc(sizeof(uint8_t) * IMG_ROW * IMG_COL * IMG_CHANEL));
    if (rawBuf == nullptr)
    {
        std::printf("[loadJpg_mmap_dec2_resz #%u] Allocating JPEG buffer FAILED\n", tid);
        return -1;
    }
    int nitem = 0;
    while (true)
    {
        message info = queue.get();
        if (info.ptr == nullptr)
        {
            break;
        }

        size_t jpegSize = zq_loader_mmap(info.path.c_str(), &jpegMap);
        zq_decoder2(handle, jpegMap, jpegSize, rawBuf);
        munmap(static_cast<void*>(jpegMap), jpegSize);
        jpegMap = nullptr;

        zq_resize(rawBuf, info.ptr);

        counter.count();
        ++nitem;
    }
    tjDestroy(handle);
    free(rawBuf);
    std::printf("[loadJpg_mmap_dec2_resz #%u] loaded %d jpgs, now exiting...\n", tid, nitem);
    return nitem;
};


int loadJpg_mmap_dec2r_rez2(const unsigned tid, zq_queue& queue, counter_mt& counter)
{
    tjhandle handle = tjInitDecompress();
    uint8_t* jpegMap = nullptr;
    uint8_t*  rawBuf = static_cast<uint8_t*>(malloc(sizeof(uint8_t) * IMG_ROW * IMG_COL * IMG_CHANEL));
    if (rawBuf == nullptr)
    {
        std::printf("[loadJpg_mmap_dec2r_rez2 #%u] Allocating JPEG buffer FAILED\n", tid);
        return -1;
    }
    int nitem = 0;
    while (true)
    {
        message info = queue.get();
        if (info.ptr == nullptr)
        {
            break;
        }

        size_t jpegSize = zq_loader_mmap(info.path.c_str(), &jpegMap);
        zq_decoder2_resz(handle, jpegMap, jpegSize, rawBuf);
        munmap(static_cast<void*>(jpegMap), jpegSize);
        jpegMap = nullptr;

        zq_resize2(rawBuf, info.ptr);

        counter.count();
        ++nitem;
    }
    tjDestroy(handle);
    free(rawBuf);
    std::printf("[loadJpg_mmap_dec2r_rez2 #%u] loaded %d jpgs, now exiting...\n", tid, nitem);
    return nitem;
};


int iotest_fread(const unsigned tid, zq_queue& queue, counter_mt& counter)
{
    int nitem = 0;
    while (true)
    {
        message info = queue.get();
        if (info.ptr == nullptr)
        {
            break;
        }

        size_t jpegSize = zq_loader(info.path.c_str(), info.ptr);

        counter.count();
        ++nitem;
    }
    std::printf("[iotest_fread #%u] loaded %d jpgs, now exiting...\n", tid, nitem);
    return nitem;
}


int iotest_mmap(const unsigned tid, zq_queue& queue, counter_mt& counter)
{
    uint8_t* jpegMap = nullptr;
    int nitem = 0;
    while (true)
    {
        message info = queue.get();
        if (info.ptr == nullptr)
        {
            break;
        }

        size_t jpegSize = zq_loader_mmap(info.path.c_str(), &jpegMap);
        memcpy(info.ptr, jpegMap, jpegSize);
        munmap(static_cast<void*>(jpegMap), jpegSize);
        jpegMap = nullptr;

        counter.count();
        ++nitem;
    }
    std::printf("[iotest_mmap #%u] loaded %d jpgs, now exiting...\n", tid, nitem);
    return nitem;
}


size_t zq_loader(const char* fname, uint8_t* dst)
{
    FILE* jpegFile = fopen(fname, "rb");
    if (jpegFile == nullptr)
    {
        std::printf("[zq_loader] Failed to open [%s]\n", fname);
        return 0;
    }
    fseek(jpegFile, 0, SEEK_END);
    ssize_t jpegSize = ftell(jpegFile);
    if (jpegSize < 0)
    {
        std::printf("[zq_loader] Failed to determine [%s] size\n", fname);
        fclose(jpegFile);
        return 0;
    }
    rewind(jpegFile);
    size_t readCnt = fread(dst, jpegSize, 1, jpegFile);
    fclose(jpegFile);
    if (readCnt < 1)
    {
        std::printf("[zq_loader] Failed to load [%s]\n", fname);
        return 0;
    }
    return size_t(jpegSize);
}

size_t zq_loader_mmap(const char* fname, uint8_t** dst)
{
    int fd = open(fname, O_RDONLY);
    if (fd < 0)
    {
        std::printf("[zq_loader_mmap] Failed to open [%s]\n", fname);
        return 0;
    }
    off_t fileLen = lseek(fd, 0, SEEK_END);
    if (fileLen < 0)
    {
        std::printf("[zq_loader_mmap] Failed to determine [%s] size\n", fname);
        close(fd);
        return 0;
    }
    size_t jpegSize = size_t(fileLen);
    uint8_t* jpegMap = static_cast<uint8_t*>(mmap(nullptr, jpegSize, PROT_READ, MAP_PRIVATE, fd, 0));
    if (jpegMap == MAP_FAILED)
    {
        std::printf("[zq_loader_mmap] Failed to map [%s] Errno: %d\n", fname, errno);
        close(fd);
        return 0;
    }
    close(fd);
    *dst = jpegMap;
    return jpegSize;
}

int zq_decoder(tjhandle handle, const uint8_t* src, size_t jpegSize, uint8_t* dst)
{
    int ret = tjDecompress2(handle, src, jpegSize, dst, IMG_COL, 0, IMG_ROW, TJPF_RGB, 0);
    if ((0 != tjGetErrorCode(handle)) && (ret < 0))
    {
        std::printf(
            "[zq_decoder] error :decompress to yuv failed, errorStr:%s, errorCode:%d\n",
            tjGetErrorStr(), tjGetErrorCode(handle)
        );
    }
    if ((0 == tjGetErrorCode(handle)) && (ret < 0))
    {
        std::printf("[zq_decoder] warning :errorStr:%s, errorCode:%d\n", tjGetErrorStr(), tjGetErrorCode(handle));
    }
    return ret;
}

int zq_decoder2(tjhandle handle, const uint8_t* src, size_t jpegSize, uint8_t* dst)
{
    int ret = tjDecompress2(handle, src, jpegSize, dst, IMG_COL, 0, IMG_ROW, TJPF_RGB, TJFLAG_FASTDCT | TJFLAG_FASTUPSAMPLE);
    if ((0 != tjGetErrorCode(handle)) && (ret < 0))
    {
        std::printf(
            "[zq_decoder2] error :decompress to yuv failed, errorStr:%s, errorCode:%d\n",
            tjGetErrorStr(), tjGetErrorCode(handle)
        );
    }
    if ((0 == tjGetErrorCode(handle)) && (ret < 0))
    {
        std::printf("[zq_decoder2] warning :errorStr:%s, errorCode:%d\n", tjGetErrorStr(), tjGetErrorCode(handle));
    }
    return ret;
}

int zq_decoder2_resz(tjhandle handle, const uint8_t* src, size_t jpegSize, uint8_t* dst)
{
    // (360, 640) --> (180, 320)
    int ret = tjDecompress2(handle, src, jpegSize, dst, IMG_COL / 2, 0, IMG_ROW / 2, TJPF_RGB, TJFLAG_FASTDCT | TJFLAG_FASTUPSAMPLE);
    if ((0 != tjGetErrorCode(handle)) && (ret < 0))
    {
        std::printf(
            "[zq_decoder2_resz] error :decompress to yuv failed, errorStr:%s, errorCode:%d\n",
            tjGetErrorStr(), tjGetErrorCode(handle)
        );
    }
    if ((0 == tjGetErrorCode(handle)) && (ret < 0))
    {
        std::printf("[zq_decoder2_resz] warning :errorStr:%s, errorCode:%d\n", tjGetErrorStr(), tjGetErrorCode(handle));
    }
    return ret;
}

void zq_resize(const uint8_t* src, uint8_t* dst)
{
    static_assert(sizeof(pixel) == 3, "zq_resize");
    const pixel* psrc = reinterpret_cast<const pixel*>(src);
    pixel* pdst = reinterpret_cast<pixel*>(dst);

    // (360, 640) --> (160, 320)
    for (unsigned r = 0; r < RESIZE_ROW; ++r)
    {
        float rr = r * float(IMG_ROW - 1) / float(RESIZE_ROW - 1) + 0.5;
        // float rr = (r + 0.5) * float(IMG_ROW) / float(RESIZE_ROW);

        const pixel* ptr = psrc + IMG_COL * unsigned(rr) + 1;
        for (unsigned c = 0; c < RESIZE_COL; ++c)
        {
            *pdst = *ptr;
            pdst++;
            ptr += 2;
        }
    }
    return;
}

void zq_resize2(const uint8_t* src, uint8_t* dst)
{
    // (180, 320) --> (160, 320)
    for (unsigned r = 0; r < RESIZE_ROW; ++r)
    {
        float rr = r * float(IMG_ROW / 2 - 1) / float(RESIZE_ROW - 1) + 0.5;
        // float rr = (r + 0.5) * float(IMG_ROW / 2) / float(RESIZE_ROW);
        const uint8_t* ptr = src + RESIZE_COL * 3 * unsigned(rr);
        memcpy(dst, ptr, RESIZE_COL * 3);
        dst += RESIZE_COL * 3;
    }
    return;
}

